import tensorflow as tf
from keras.preprocessing.image import ImageDataGenerator
from keras.layers import Dense
from keras.models import Sequential
from keras.layers import Conv2D, MaxPooling2D, Flatten
# 导入优化器
from keras.optimizers import Adam
import os
import shutil

# 定义图像的宽度和高度
image_width, image_height = 224, 224
# 定义训练和验证数据的路径
batch_size = 32
# 定义训练和验证数据的路径
train_data_dir = './training_set'
test_data_dir = './test_set'


def create_model():
    # 创建训练数据生成器并应用数据增强
    train_datagn = ImageDataGenerator(
        rescale=1. / 255,  # 缩放
        shear_range=0.2,  # 剪切范围
        zoom_range=0.2,  # 缩放范围 1
        horizontal_flip=True  # 水平旋转
    )

    # 创建验证数据生成器并应用数据增强
    test_datagen = ImageDataGenerator(
        rescale=1. / 255  # 缩放范围0-1
    )

    # 加载训练图像
    train_generator = train_datagn.flow_from_directory(
        train_data_dir,
        target_size=(image_width, image_height),  # 调整图像大小
        batch_size=batch_size,  # 批次大小
        class_mode='categorical'  # 分类
    )

    # 构建CNN网络
    model = Sequential([
        Conv2D(32, (3, 3), input_shape=(224, 224, 3), activation='relu'),
        MaxPooling2D((2, 2)),

        Conv2D(64, (3, 3), activation='relu'),
        MaxPooling2D((2, 2)),

        Flatten(),
        Dense(128, activation='relu'),
        #添加输出层，神经元的数量等于训练数据的类别数量，激活函数为Softmax，用于输出每个类别的概率
        Dense(units=len(train_generator.class_indices), activation='softmax')
    ])

    # 编译模型
    model.compile(
        #自适应学习率的优化算法(随机梯度下降)
        optimizer='adam',
        #交叉熵作为损失函数，它是多分类问题的常用损失函数
        loss='categorical_crossentropy',
        # 使用准确率作为评估指标
        metrics=['accuracy']
    )

    # 训练模型
    history = model.fit(
        train_generator,
        epochs=40,
        # 每个周期的步数等于训练集的样本数量除以批量大小
        steps_per_epoch=len(train_generator)
    )

    # 模型保存
    model.save('model2.h5')


# def test_model():
#     model = tf.keras.models.load_model('./model.h5')
#     # 测试集预测
#     test_folder = './test_set/'
#     bird_classes = [
#         'Asian Green Bee-Eater(食蜂鸟)',
#         'Brown-Headed Barbet(斑头绿拟啄木鸟)',
#         'Cattle Egret(牛背鹭)',
#         'Common Kingfisher(普通翠鸟)',
#         'Common Myna(新西兰八哥)',
#         'Common Rosefinch(普通朱雀)',
#         'Common Tailorbird(普通缝叶莺)',
#         'Coppersmith Barbet(赤胸拟啄木鸟)',
#         'Forest Wagtail(山鹡鸰)',
#         'Gray Wagtail(灰鹡鸰)',
#         'Hoopoe(戴胜)',
#         'House Crow(家鸦)',
#         'Indian Grey Hornbill(印度灰犀鸟)',
#         'Indian Peacock(印度孔雀)',
#         'Indian Pitta(蓝翅八色鸫)',
#         'Indian Roller(棕胸佛法僧)',
#         'Jungle Babbler(丛林鸫鹛)',
#         'Northern Lapwing(凤头麦鸡)',
#         'Red-Wattled Lapwing(肉垂麦鸡)',
#         'Ruddy Shelduck(赤麻鸭)',
#         'Rufous Treepie(棕腹树鹊)',
#         'Sarus Crane(赤颈鹤)',
#         'White-Breasted Kingfisher(白胸翠鸟)',
#         'White-Breasted Waterhen(白胸苦恶鸟)',
#         'White Wagtail(白鹡鸰)'
#     ]
#
#     #创建图片保存文件判断文件是否存在，不存在则创建
#     test_result_flg = './test_result/'
#     subfolder_names = [
#         'Asian Green Bee-Eater',
#         'Brown-Headed Barbet',
#         'Cattle Egret',
#         'Common Kingfisher',
#         'Common Myna',
#         'Common Rosefinch',
#         'Common Tailorbird',
#         'Coppersmith Barbet',
#         'Forest Wagtail',
#         'Gray Wagtail',
#         'Hoopoe',
#         'House Crow',
#         'Indian Grey Hornbill',
#         'Indian Peacock',
#         'Indian Pitta',
#         'Indian Roller',
#         'Jungle Babbler',
#         'Northern Lapwing',
#         'Red-Wattled Lapwing',
#         'Ruddy Shelduck',
#         'Rufous Treepie',
#         'Sarus Crane',
#         'White-Breasted Kingfisher',
#         'White-Breasted Waterhen',
#         'White Wagtail'
#     ]
#     #循环判断
#     for i in subfolder_names:
#         path = os.path.join(test_result_flg, i)
#         print(path)
#         if not os.path.isdir(path):
#             os.mkdir(path)
#         else:
#             print("文件夹已存在！")
#
#     #测试图片名字集合
#     test_images = os.listdir(test_folder)
#     #循环遍历判断每一张图片的相似度然后保存到结果图片的路径中
#     for img_name in test_images:
#         img_path = os.path.join(test_folder, img_name)
#         img = tf.keras.preprocessing.image.load_img(img_path, target_size=(image_width, image_height))
#         img_array = tf.keras.preprocessing.image.img_to_array(img)
#         img_array = img_array / 255.
#         img_array = img_array.reshape(1, *img_array.shape)
#         # 输出预测结果
#         prediction = model.predict(img_array)[0]
#         # 将预测结果和鸟类名字对应
#         for bird, score in zip(bird_classes, prediction):
#             print(f'{bird}:{score * 100:.2f}%')
#         # 显示最有可能的鸟类
#         most_likely_bird = bird_classes[prediction.argmax()]
#         #分割出鸟类的名字依次放入对应名字的文件夹
#         bird_flg_name = most_likely_bird.split('(')[0]
#         print(bird_flg_name)
#         #移动图片到指定的文件夹
#         # test1 = test_folder+img_name
#         # print(test1)
#         # test2 = test_result_flg+bird_flg_name+'/'+img_name
#         # print(test2)
#         shutil.move(test_folder+img_name,test_result_flg+bird_flg_name+'/'+img_name)
#         print(f'\n{img_name} 最高相似度是---> {most_likely_bird}\n')
#         print('-' * 100)
#
# def return_result(imgae_name):
#     #模型加载
#     model = tf.keras.models.load_model('./model.h5')
#     #鸟类分类
#     bird_classes = [
#         'Asian Green Bee-Eater(食蜂鸟)',
#         'Brown-Headed Barbet(斑头绿拟啄木鸟)',
#         'Cattle Egret(牛背鹭)',
#         'Common Kingfisher(普通翠鸟)',
#         'Common Myna(新西兰八哥)',
#         'Common Rosefinch(普通朱雀)',
#         'Common Tailorbird(普通缝叶莺)',
#         'Coppersmith Barbet(赤胸拟啄木鸟)',
#         'Forest Wagtail(山鹡鸰)',
#         'Gray Wagtail(灰鹡鸰)',
#         'Hoopoe(戴胜)',
#         'House Crow(家鸦)',
#         'Indian Grey Hornbill(印度灰犀鸟)',
#         'Indian Peacock(印度孔雀)',
#         'Indian Pitta(蓝翅八色鸫)',
#         'Indian Roller(棕胸佛法僧)',
#         'Jungle Babbler(丛林鸫鹛)',
#         'Northern Lapwing(凤头麦鸡)',
#         'Red-Wattled Lapwing(肉垂麦鸡)',
#         'Ruddy Shelduck(赤麻鸭)',
#         'Rufous Treepie(棕腹树鹊)',
#         'Sarus Crane(赤颈鹤)',
#         'White-Breasted Kingfisher(白胸翠鸟)',
#         'White-Breasted Waterhen(白胸苦恶鸟)',
#         'White Wagtail(白鹡鸰)'
#     ]
#     #用户上传文件夹路径
#     user_path = './Internet_birds_images/'
#     #对图片数据进行预处理
#     img_path = os.path.join(user_path, imgae_name)
#     img = tf.keras.preprocessing.image.load_img(img_path, target_size=(image_width, image_height))
#     img_array = tf.keras.preprocessing.image.img_to_array(img)
#     img_array = img_array / 255.
#     img_array = img_array.reshape(1, *img_array.shape)
#     # 输出预测结果
#     prediction = model.predict(img_array)[0]
#     # 将预测结果和鸟类名字对应
#     for bird, score in zip(bird_classes, prediction):
#         print(f'{bird}:{score * 100:.2f}%')
#     # 显示最有可能的鸟类
#     most_likely_bird = bird_classes[prediction.argmax()]
#     # 分割出鸟类的名字依次放入对应名字的文件夹
#     bird_flg_name = most_likely_bird.split('(')[0]
#     print(bird_flg_name)
#     return 0

if __name__ == '__main__':
    print(1)
    create_model()